import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, RandomSampler
import matplotlib.pyplot as plt
import numpy as np
import torchvision
from functools import partial
# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5,), std=(0.5,))  # 将图像归一化到[-1, 1]范围内
])

# 下载并加载训练数据
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=False)
sample_index = [i*10 for i in range(6000)] #假设取前500个训练数据
X_train = []
y_train = []
for i in sample_index:
    X = train_dataset[i][0]
    X_train.append(X)
    y = train_dataset[i][1]
    y_train.append(y)

sampled_train_data = [(X, y) for X, y in zip(X_train, y_train)] #包装为数据对
trainDataLoader = torch.utils.data.DataLoader(sampled_train_data, batch_size=16, shuffle=True)

train_loader = DataLoader(dataset=sampled_train_data, batch_size=256, shuffle=True)

class Generator(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, output_size),
            nn.Tanh()
        )

    def forward(self, x):
        return self.main(x)


class Discriminator(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_size, hidden_size),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_size, output_size),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.main(x)

# 定义模型参数
input_size = 100  # 噪声向量的维度
hidden_size = 256
image_size = 28 * 28  # MNIST图像的维度

# 创建生成器和判别器实例
G = Generator(input_size, hidden_size, image_size)
D = Discriminator(image_size, hidden_size, 1)

def load_models(G, D, G_path='generator.pth', D_path='discriminator.pth'):
    """
    加载保存的生成器和判别器模型

    参数:
        G: 生成器模型实例(未训练)
        D: 判别器模型实例(未训练)
        G_path: 生成器模型路径
        D_path: 判别器模型路径

    返回:
        加载好参数的生成器和判别器
    """
    # 加载生成器
    G.load_state_dict(torch.load(G_path))
    G.eval()  # 设置为评估模式

    # 加载判别器
    D.load_state_dict(torch.load(D_path))
    D.eval()  # 设置为评估模式

    print(f"模型已从 {G_path} 和 {D_path} 加载")

    return G, D
G,D = load_models(G,D)
# 生成一些新图像
z = torch.randn(64, input_size)
fake_images = G(z)
fake_images = fake_images.view(fake_images.size(0), 1, 28, 28)

# 可视化生成的图像
grid = torchvision.utils.make_grid(fake_images, nrow=8, normalize=True)
plt.imshow(grid.permute(1, 2, 0).detach().numpy())
plt.title('Generated Images')
plt.show()


